#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import os
import sys
import argparse

from elftools.elf.elffile import ELFFile
from elftools.elf.descriptions import describe_sh_flags

from prettytable import PrettyTable, MARKDOWN
from ctools.common.myhtmltable import MyHTMLTable
from ctools.gvhelper.graphwrapper import GraphWrapper
from ctools.subcmd.ctbase import CToolsCmd

import logging
from ctools.logext import *

_fpath = os.path.basename(os.path.realpath(__file__))
_rlog, _plog, _mlog, _tlog = get_loggers(__file__, __name__, propagate=False)

def _format_hex(val:int, nbits=0) -> str:
    masks = (0xFFFF000000000000, 0x0000FFFF00000000,
             0x00000000FFFF0000, 0x000000000000FFFF)
    _n = 0
    for idx, mask in enumerate(masks):
        if val & mask:
            _n = (len(masks) - idx) * 16
            break
    if nbits > _n:
        _n = nbits
    ret = []
    for idx in range(0, _n, 16):
        ret.insert(0, '{:04X}'.format((val >> idx) & 0xFFFF))
    return '_'.join(ret)

class ElfFormater(object):
    def __init__(self):
        '''
        Elf64_Word sh_name;		/* Section name, index in string tbl */
        Elf64_Word sh_type;		/* Type of section */
        Elf64_Xword sh_flags;		/* Miscellaneous section attributes */
        Elf64_Addr sh_addr;		/* Section virtual addr at execution */
        Elf64_Off sh_offset;		/* Section file offset */
        Elf64_Xword sh_size;		/* Size of section in bytes */
        Elf64_Word sh_link;		/* Index of another section */
        Elf64_Word sh_info;		/* Additional section information */
        Elf64_Xword sh_addralign;	/* Section alignment */
        Elf64_Xword sh_entsize;	/* Entry size if section holds table */
        '''

        self._shdr_fields = (
            "sh_name", "sh_type", "sh_flags",
            "sh_addr", "sh_offset",
            "sh_size", "sh_link",
            "sh_info", "sh_addralign", "sh_entsize",
        )

    def _get_field(self, section, field, raw=False):
        if raw:
            return section[field]
        if field == 'sh_name':
            return section.name
        if field in ("sh_flags",):
            return describe_sh_flags(section[field])
        return section[field]

    def _get_info(self, elffile, fields:list|tuple):
        descs = []
        with open(elffile, "rb") as efd:
            elf = ELFFile(efd)
            for idx, section in enumerate(elf.iter_sections()):
                if section['sh_addr'] == 0:
                    continue
                desc = [self._get_field(section, field) for field in fields]
                descs.append(desc)
        return descs
'''
markdown:
|   name  | flags | addr start |  addr end |  size  |
| :-----: | :---: | :--------: | :-------: | :----: |
TODO:
| name | flags | addr | offset | size | link | info | align | entsize |
| ---- | ----- | ---- | ------ | ---- | ---- | ---- | ----- | ------- |
'''
class ElfMDTableFormater(ElfFormater):
    def _process_desc_holes(self, descs:list):
        _descs = [descs[0], ]
        # add hole line
        for idx, l in enumerate(descs[1:]):
            addr1, size1 = descs[idx][-2:]
            addr2, size2 = l[-2:]
            if addr1 + size1 != addr2:
                _l = ('hole', '', addr1 + size1, addr2 - (addr1 + size1))
                _descs.append(_l)
            _descs.append(l)
        return _descs

    def _addrow_from_desc(self, table:PrettyTable, desc):
        name, flags, addr, size = desc
        row = (
            name, flags,
            _format_hex(addr), _format_hex(addr + size - 1),
            f'0x{size:X}({size:d})'
        )
        table.add_row(row)

    def process(self, elffile:str, fields=None):
        # FIXME: fields
        fields = ('sh_name', 'sh_flags', 'sh_addr', 'sh_size')

        descs = self._get_info(elffile, fields)

        # sort by sh_addr
        index = fields.index('sh_addr')
        descs = sorted(descs, key=lambda x: x[index])

        lines = self._process_desc_holes(descs)

        field_names = ('name', 'flags', 'addr start', 'addr end', 'size',)
        table = PrettyTable(field_names=field_names)
        table.align['addr start'] = 'l'
        table.align['addr end'] = 'l'

        for idx, l in enumerate(lines):
            self._addrow_from_desc(table, l)

        table.set_style(MARKDOWN)
        print(table.get_string())

'''
graphviz: html/svg/png/pdf
|------------------|------------------|
|                  |   high addr      |
|                  |                  |
|   name           |                  |
|   flags          |     size         |
|                  |                  |
|                  |                  |
|                  |    low addr      |
|------------------|------------------|
|     hole         |    size          |
|------------------|------------------|
'''
class ElfGvFormater(ElfFormater):

    def _process_desc_holes(self, descs:list):
        _descs = [descs[0], ]
        # add hole line
        for idx, l in enumerate(descs[1:]):
            addr1, size1 = descs[idx][-2:]
            addr2, size2 = l[-2:]
            if addr1 + size1 != addr2:
                _l = ('hole', '', addr1 + size1, addr2 - (addr1 + size1))
                _descs.append(_l)
            _descs.append(l)
        return _descs

    def _addrow_from_desc_flat(self, htable:MyHTMLTable, desc):
        name, flags, addr, size = desc

        row = (
            name, flags,
            _format_hex(addr), _format_hex(addr + size - 1),
            f'0x{size:X}({size:d})'
        )
        htable.add_row(row)
        if name.startswith('hole'):
            return

        portname = name.strip('.').replace('.', '_')
        htable[-1][0].attr.PORT = portname

    def _addrow_from_desc_pile(self, htable:MyHTMLTable, desc):
        name, flags, addr, size = desc

        name_flags = name
        if flags:
            name_flags += '\n' + flags

        if name.startswith('hole'):
            addr_info = "0x{:X}({:d})".format(size, size)
        else:
            addr_info = "{}\n\n0x{:X}({:d})\n\n{}".format(
                            _format_hex(addr + size - 1),
                            size, size,
                            _format_hex(addr)
                        )
        row = (name_flags, addr_info,)
        htable.add_row(row)
        if name.startswith('hole'):
            return

        portname = name.strip('.').replace('.', '_')
        htable[-1][0].attr.PORT = portname
        htable[-1][1].attr.PORT = '{}_addr'.format(portname)

    def _process_graph(self, elffile:str, style:str, fields:list|tuple=None):
        func = self._addrow_from_desc_flat
        if style == 'pile':
            func = self._addrow_from_desc_pile

        descs = self._get_info(elffile, fields)

        # sort by sh_addr
        index = fields.index('sh_addr')
        descs = sorted(descs, key=lambda x: x[index])

        lines = self._process_desc_holes(descs)

        htable = MyHTMLTable()
        htable.attr.border = 0
        htable.attr.cellborder = 1
        htable.attr.cellspacing = 0

        bgcolors = ("gray", "white")
        # put high addr to top
        lines.reverse()
        for idx, l in enumerate(lines):
            func(htable, l)
            for cell in htable[-1]:
                cell.attr.bgcolor = bgcolors[idx % 2]

        return htable.to_html()

    def process_graph(self, elffile:str, outfile:str, style:str, format:str=None, fields:list|tuple=None):
        # FIXME: fields
        fields = ('sh_name', 'sh_flags', 'sh_addr', 'sh_size')
        html = self._process_graph(elffile, style, fields)
        gvtext = 'digraph G {\n'
        gvtext += f'n [shape=plaintext, label=<{html}>]\n'
        gvtext += '}\n'

        graph = GraphWrapper(string=gvtext)
        graph.draw(outfile, format=format, prog='dot')

class ElfLayoutCmd(CToolsCmd):
    @classmethod
    def name(cls):
        return "elflayout"

    @classmethod
    def description(cls):
        return "elf layout tool"

    def do_command(self, argv:list[str]):
        argp = argparse.ArgumentParser(description='ELF view tool')
        argp.add_argument("-e", "--engine", type=str, choices=("md", "gv", "dot"), default="gv")
        argp.add_argument("-o", "--output", type=str, help="required for gv/dot")
        argp.add_argument("-s", "--style", type=str, choices=('flat', 'pile'), default='flat',
                            help="required for gv/dot")
        argp.add_argument("-f", "--format", type=str, default=None,
                            help="required for gv/dot")
        argp.add_argument('file', type=str)

        args = argp.parse_args(argv)

        if args.engine in ("gv", "dot"):
            if args.output is None:
                argp.print_usage()
                sys.exit(-1)

            gvf = ElfGvFormater()
            gvf.process_graph(args.file, args.output, args.style, args.format)
        elif args.engine == "md":
            mdf = ElfMDTableFormater()
            mdf.process(args.file, args.output)

if __name__ == '__main__':
    ElfLayoutCmd().do_command(sys.argv[1:])
